import base64
import datetime
import json
import uuid

import requests
from loguru import logger

from src.config.redis_cache import RedisCache
from src.model.scan_task import ScanTask
from src.model.scan_work import ScanWork
from src.service.web_scraper_service import WebScraperService
from src.utils.port_scan import PortScanner
from src.utils.read_config import server_address_config

remote_url, is_proactive_request,file_path = server_address_config()

REMOTE_URL  = remote_url
IS_PROACTIVE_REQUEST  = is_proactive_request
FILE_PATH  = file_path



class SchedulerJob:
    def __init__(self):
        self.redis_cache = RedisCache()

    # 定时去对IP的端口进行扫描，如果扫描到端口就进行保存到redis中，然后进行回调。
    def process_tasks(self, app):
        with app.app_context():
            # 从redis中获取是否已经有执行完的任务，如果有执行完的任务获取到结果然后进行回调后再处理下一个任务的执行。
            ports_result_keys = self.redis_cache.keys_by_pattern("scan:ports:result*")  # 使用实例属性
            pending_task = None
            flag = True
            if ports_result_keys:
                # 处理结果并进行回调
                scan_result = {}
                for result_key in ports_result_keys:
                    result = self.redis_cache.get(result_key)
                    uuid = result_key.split(":")[3]
                    if result:
                        scan_result[result_key] = result
                    # 查询这个子任务下面是否还有其他子任务在执行没，如果没有的话进行更新主任务状态需要进行执行下一步，并且进行内容回调
                    scan_task = ScanTask()
                    scan_task.uuid = uuid
                    scan_task = scan_task.query_by_uuid(uuid)
                    work_uuid = scan_task.work_uuid;
                    scan_task.status = "execution"
                    scan_task_tmp = scan_task.get_work_uuid_task_status(work_uuid, scan_task.status);
                    if scan_task_tmp:
                        continue
                    scan_task.status = "pending"
                    scan_task = scan_task.get_work_uuid_task_status(work_uuid, scan_task.status);
                    if not scan_task:
                        scan_work = ScanWork()
                        scan_work.uuid = work_uuid
                        scan_work = scan_work.get_uuid_work(work_uuid)
                        scan_work.status = "done"
                        scan_work.callback_status = "done"
                        scan_work.update()
                        self.callback_work(scan_work)
                        # 可以选择删除已处理的结果键，避免重复处理
                        self.redis_cache.remove(result_key)
                    else:
                        if flag:
                            flag = False
                            pending_task = scan_task
                # 不管是否有完成的直接进行下一步
                if not flag and pending_task:
                    self.next_task_step(pending_task)
            else:
                pending_task = ScanTask()
                pending_task.status = "execution"
                pending_task = pending_task.get_status_task(pending_task.status)
                if pending_task:
                    return
                pending_task = ScanTask()
                pending_task.status = "pending"
                pending_task = pending_task.get_status_task(pending_task.status)
                if pending_task:
                    self.next_task_step(pending_task)

    def callback_work(self, scan_word: ScanWork):
        if scan_word:
            print(f"Processing result: {scan_word}")
            # 获取该任务下的所有子任务列表成功的列表
            scan_task = ScanTask()
            scan_task_list = scan_task.get_work_uuid_task(scan_word.uuid)
            if scan_task_list:
                scan_task_ip = {}
                dead_hosts_all = []
                for scan_task in scan_task_list:
                    if scan_task.status == "done":
                        scan_task_list_key = self.redis_cache.keys_by_pattern(
                            'scan:ports:single:' + scan_task.uuid + ":*")
                        scan_task_list_key.sort()
                        task_uuids = '';
                        for port_result_key in scan_task_list_key:
                            port_result = self.redis_cache.get(port_result_key)
                            if port_result:
                                ip = port_result_key.replace('scan:ports:single:' + scan_task.uuid + ":", '');
                                scan_task_ip[ip] = port_result
                            self.redis_cache.remove(port_result_key)
                            task_uuids = scan_task.uuid;
                        # 可以选择删除已处理的结果键，避免重复处理
                        dead_hosts = self.redis_cache.lrange("scan:ports:dead_hosts:" + task_uuids)
                        if dead_hosts:
                            dead_hosts_all.extend(dead_hosts)
                        self.redis_cache.remove("scan:ports:dead_hosts:" + scan_task.uuid)
                if scan_task_ip:
                    headers = {
                        "Content-Type": "application/json"
                    }

                    data = {
                        'alive_ips': json.dumps(scan_task_ip),
                        'probe_name': '',
                        'inactive_ips': json.dumps(dead_hosts_all)

                    }
                    try:
                        response = requests.post(scan_word.callback_url, json=data, headers=headers)
                    except requests.exceptions.RequestException as e:
                        print(f"Error sending request: {e}")

    def next_task_step(self, scan_task: ScanTask):
        scanner = PortScanner()
        if scan_task and scan_task.status == "pending":
            scan_task.status = "execution"
            scan_task.update()
            target = scan_task.params["ip"]
            # 判断 target 的类型并进行相应的处理
            if isinstance(target, list):
                target = ",".join(target)
            ports = scan_task.params['port']
            uuid = scan_task.uuid
            id = scan_task.id
            scan_task.closeSession()
            scanner.scan_and_discover(target, ports, 1000, uuid)
            scan_task_new = ScanTask()
            scan_task_new = scan_task_new.get_pending_task("execution", uuid)
            scan_task_new.status = "done"
            scan_task_new.update()

    # 定时去获取需要扫描的IP信息进行存储入库
    def get_ip_info(self, app):
        with app.app_context():
            headers = {
                "Content-Type": "application/json"
            }
            data = {
                'probe_name': ''
            }
            try:
                response = requests.post(REMOTE_URL+"/ip/scan", json=data, headers=headers)
                # 处理响应数据 从中获取到需要扫描的IP信息并进行入库操作。
                if response.status_code == 200:
                    # 处理响应数据
                    params_str = response.json.get("params")
                    execute_time_str = response.json.get("execute_time")
                    # 将params字符串解析为字典
                    if isinstance(params_str, str):
                        params = json.loads(params_str)
                    else:
                        params = params_str

                    if isinstance(execute_time_str, str):
                        execute_time = json.loads(execute_time_str)
                    else:
                        execute_time = execute_time_str

                    callback_url = response.json.get("callback_url")  # 修改变量名为小写加下划线
                    progress_url = response.json.get("progress_url", "http")

                    logger.info(f'params:{params},callback_url:{callback_url}')  # 修改日志输出中的变量名

                    # 创建scanwork对象并对其中的值赋值
                    data = {
                        "params": params,
                        "callback_url": callback_url,
                        "callback_type": 'API',
                        "callback_status": 'pending',
                        "execute_time": execute_time,
                        "uuid": str(uuid.uuid4()),  # 假设使用uuid库生成唯一标识符
                        "status": 'pending',  # 假设初始状态为'PENDING'
                        "source": 'default',  # 根据实际情况设置默认值或从其他地方获取
                        "work_type": 'portscan',  # 根据实际情况设置默认值或从其他地方获取
                        "retry": 0,  # 假设初始重试次数为0
                        "priority": 1,  # 假设默认优先级为1
                        "progress": 0,  # 假设初始进度为0
                        "progress_type": 'API',  # 根据实际情况设置默认值或从其他地方获取
                        "progress_url": progress_url,  # 根据实际情况设置默认值或从其他地方获取
                        "create_time": datetime.datetime.now(),  # 使用当前时间作为创建时间
                        "update_time": datetime.datetime.now()  # 使用当前时间作为更新时间
                    }
                    new_data = ScanWork(**data)
                    # 创建scanwork对象并对其中的值赋值
                    scan_work = ScanWork.insert(new_data)
                    if not scan_work:
                        logger.error("scanwork insert error")
                    # 拆分扫描任务，如果包含的IP太多的话进行任务拆分
                    ips = params.get('ip', '').split(',')
                    port_str = params.get('port', '')
                    num = 256
                    if len(ips) > num:
                        for i in range(0, len(ips), num):
                            task_data = {
                                "execute_time": execute_time,
                                "work_uuid": scan_work.uuid,  # 假设使用uuid库生成唯一标识符
                                "uuid": str(uuid.uuid4()),
                                "status": 'pending',  # 假设初始状态为'PENDING'
                                "retry": 1,
                                "priority": 1,
                                "create_time": datetime.datetime.now(),  # 使用当前时间作为创建时间
                                "update_time": datetime.datetime.now()  # 使用当前时间作为更新时间
                            }
                            task_data['params'] = {
                                "ip": ips[i:i + num],
                                "port": port_str
                            }
                            new_ips_data = ScanTask(**task_data)
                            ScanTask.insert(new_ips_data)
                    else:
                        task_data = {
                            "execute_time": execute_time,
                            "work_uuid": str(scan_work.uuid),  # 假设使用uuid库生成唯一标识符
                            "uuid": str(uuid.uuid4()),
                            "status": 'pending',  # 假设初始状态为'PENDING'
                            "retry": 1,
                            "priority": 1,
                            "create_time": datetime.datetime.now(),  # 使用当前时间作为创建时间
                            "update_time": datetime.datetime.now()  # 使用当前时间作为更新时间
                        }
                        task_data['params'] = {
                            "ip": ips,
                            "port": port_str
                        }
                        new_ips_data = ScanTask(**task_data)
                        ScanTask.insert(new_ips_data)
                else:
                    logger.error(f"Request failed with status code: {response.status_code}")
            except requests.exceptions.RequestException as e:
                logger.error(f"Error sending request: {e}")

    # 定时获取网站的信息并存储到redis中
    def get_website_info(self, app):
        with app.app_context():
            logger.info("get_website_info")
            headers = {
                "Content-Type": "application/json"
            }
            data = {
                'probe_name': ''
            }
            try:
                response = requests.post(REMOTE_URL +"/web/capture", json=data, headers=headers)
                # 处理响应数据 从中获取到需要扫描的IP信息并进行入库操作。
                if response.status_code == 200:
                    response_data = response.json()
                    urls = response_data.get('urls', [])
                    if urls:
                        self.redis_cache.lpush('website:info:pending', urls);

            except requests.exceptions.RequestException as e:
                logger.error(f"Error sending request: {e}")

    # 定时获取网站的信息并存储到redis中然后进行发送改。
    def get_website_info_send(self, app):
        with app.app_context():
            urls = self.redis_cache.lrange('website:info:pending')
            if urls:
                self.redis_cache.remove('website:info:pending')
                # 让url根据每五个url进行一次拆分发送改。
                headers = {
                    "Content-Type": "application/json"
                }
                for i in range(0, len(urls), 5):
                    url_list = urls[i:i + 5]
                    service = WebScraperService(FILE_PATH)
                    results = service.run_scrape_multiple(urls)
                    # 在判断结果是否为空，如果不为空的话获取到相应的信息，判断相应的code是否为0，不为0的话获取到截图地址并将其转换为base64编码，并再方法这个json中，key为base64image，然后再请求接口/web/callback 将这五个一起进行发送。
                    if results:
                        try:
                            for result in results:
                                if result.get('code') == 0:
                                    base64_image = base64.b64encode(result.get('screenshot').read()).decode('utf-8')
                                    result.set('base64image', base64_image)

                            response = requests.post(REMOTE_URL+"/web/callback", json=results, headers=headers)
                        except requests.exceptions.RequestException as e:
                            logger.error("send error")
